
from petsc4py import PETSc


PETSc._initialize()
# 创建一个矩阵 A 和一个大小为 N 的向量 b
N = 5  # 向量的大小
A = PETSc.Mat().createAIJ(size=(N, N), nnz=3)  # 创建一个稀疏矩阵

# 创建向量 b
b = PETSc.Vec().create()  # 创建一个空向量
b.setSizes(N)  # 设置向量的大小
b.setFromOptions()  # 进行其他的设置

# 创建未知数向量 x
x = PETSc.Vec().create()  # 创建一个空向量
x.setSizes(N)  # 设置大小
x.setFromOptions()  # 进行其他的设置

# 填充矩阵 A 和向量 b
for i in range(N):
    A.setValue(i, i, 2.0)  # 对角线
    if i > 0:
        A.setValue(i, i - 1, -1.0)  # 下三角
    if i < N - 1:
        A.setValue(i, i + 1, -1.0)  # 上三角
    b.setValue(i, N + 1)  # 向量 b, 比如设置为 N + 1

# 完成矩阵和向量的构建
A.assemble()
b.assemble()

# 创建线性求解器
solver = PETSc.KSP().create()
solver.setOperators(A)

# 设置解法器选项（例如，使用 GMRES + 预条件器）
solver.setType("gmres")
pc = solver.getPC()
pc.setType("jacobi")

# 求解线性方程组 Ax = b
solver.setFromOptions()
solver.solve(b, x)

# 输出结果
print("解向量 x:", x.array)
# 结束 PETSc
PETSc._finalize()
